
/**
 * textcipher.cpp
 * 
 * TextCipher is a AES based text cipher
 * 
 * Copyright (c) 2008 Harunobu Oyama <nobukichi at gmail.com>
 * Licensed under GNU GPL version 3
 *
 */

#include <vector>
#include "AES.h"
#include "textcipher.h"

const int BLOCK_SIZE_IN_BITS = 128;
const int KEY_SIZE_IN_BITS = 256;
const int BLOCK_SIZE_IN_BYTES = BLOCK_SIZE_IN_BITS / 8;
const int KEY_SIZE_IN_BYTES = KEY_SIZE_IN_BITS / 8;

inline unsigned long getBlockCount(const QByteArray& data, int blockSize)
{
	return ((data.size() - 1) / blockSize) + 1;
}

inline void resizeByteArrayWithPadding(QByteArray& byteArray, int newSize, unsigned char padding)
{
	int orgSize = byteArray.size();
	byteArray.resize(newSize);
	for (int i = orgSize; i < newSize; ++i) {
		byteArray[i] = padding;
	}
}

inline bool isBase64(const QString& str)
{
	if (str.size() % 4 != 0) {
		return false; //! base64 string is 4 x n length
	}
	
	int paddingCount = 0; //! up to 3 padding can be at the end of string 
	for (int i = 0; i < str.size(); ++i) {
		QChar c = str.at(i);
		if (!c.isLetterOrNumber() && c != '+' && c != '/' && c != '=') {
			//! found a character invalid in base64
			return false;
		}
		if (c == '=') {
			++paddingCount;
			if (3 < paddingCount) {
				return false; //! max padding count is 3
			}
		} else if (0 < paddingCount) {
			return false; //! padding was not at the end of the string
		}
	}
	
	return true;
}

QString TextCipher::encode(const QString& plain, const QString& seed)
{
	if (plain.isEmpty()) {
		return "";
	}

	//! prepare input
	QByteArray in = plain.toUtf8();
	unsigned long blockCount = getBlockCount(in, BLOCK_SIZE_IN_BYTES);
	int dataSize = BLOCK_SIZE_IN_BYTES * blockCount;
	resizeByteArrayWithPadding(in, dataSize, 0);

	//! prepare key
	QByteArray key = seed.toUtf8();
	resizeByteArrayWithPadding(key, KEY_SIZE_IN_BYTES, key[0]);

	//! prepare output buffer
	std::vector<unsigned char> out;
	out.resize(dataSize);

	AES cipher;
	cipher.SetParameters(KEY_SIZE_IN_BITS, BLOCK_SIZE_IN_BITS);
	cipher.StartEncryption((const unsigned char*) key.data());
	cipher.Encrypt((const unsigned char*) in.data(), &out[0], blockCount);

	QByteArray result((const char*) &out[0], dataSize);
	return QString(result.toBase64());
}


QString TextCipher::decode(const QString& code, const QString& seed)
{
	if (code.isEmpty()) {
		return "";
	}

	//! check if the input string is in base64 format
	if (!isBase64(code)) {
		return code;
	}
	QByteArray in = QByteArray::fromBase64(code.toAscii());

	//! check if the data size is a multiple of BLOCK_SIZE
	int dataSize = in.size();
	if (dataSize % BLOCK_SIZE_IN_BYTES != 0) {
		return code;
	}
	
	//! prepare key
	QByteArray key = seed.toUtf8();
	resizeByteArrayWithPadding(key, KEY_SIZE_IN_BYTES, key[0]);
	
	//! prepare output buffer
	std::vector<unsigned char> out;
	out.resize(dataSize);
	
	AES cipher;
	cipher.SetParameters(KEY_SIZE_IN_BITS, BLOCK_SIZE_IN_BITS);
	cipher.StartDecryption((const unsigned char*) key.data());
	cipher.Decrypt((const unsigned char*) in.data(), &out[0], dataSize / BLOCK_SIZE_IN_BYTES);
	
	QByteArray result((const char*) &out[0], dataSize);
	return QString::fromUtf8(result.data());
}

